from BinarySearchTree import TreeNode
from BinarySearchTree import BinarySearchTree


# AVL平衡二叉搜索树，这个东西尽量实现树的平衡，因此在树节点上增加一个balanceFactor数值，这边用到了继承的东西
class balanceBinarySearchTree(BinarySearchTree):
    def __init__(self):
        BinarySearchTree.__init__(self)

    def _put(self,key,val,currentNode): # 添加元素最后一定是会添加到叶子结点的，所以会一直递归到叶子节点
        # print(1111)
        if key < currentNode.key:   # 添加到左子树
            if currentNode.hasLeftChild():
                self._put(key, val, currentNode.leftChild)  # 递归添加
            else:
                currentNode.leftChild = TreeNode(key, val, parent=currentNode)  #  最后加在了叶子节点上
                self.updateBalance(currentNode.leftChild)   # 加完之后需要更新所有父节点的balance值
        elif key > currentNode.key:
            if currentNode.hasRightChild():
                self._put(key,val,currentNode.rightChild)
            else:
                currentNode.rightChild = TreeNode(key,val,parent=currentNode)   # 没有右子树创建右子树
                print('jia4')
                self.updateBalance(currentNode.rightChild)  # 更新
        else:   # 不添加，修改值
            currentNode.payload = val

    def updateBalance(self,node):   # 更新所有父节点的balance值
        # print(node.isRightChild())
        # print(node.parent != None)
        if node.balanceFactor > 1 or node.balanceFactor < -1:
            self.rebalance(node)    # 二叉树失衡，更新二叉树
            return
        if node.parent != None:
            if node.isLeftChild():
                node.parent.balanceFactor += 1
            elif node.isRightChild():
                node.parent.balanceFactor -= 1
                # print(node.parent.balanceFactor)
            if node.parent.balanceFactor != 0:  # 父节点为0时说明原本为1/-1变成0，或者本来就是零，这样其父节点的父节点balance不变，因为没有增加深度
                self.updateBalance(node.parent)

    # 失衡时的旋转，试验一下就知道失衡往往都是根节点失衡，所以都是根节点旋转
    def rebalance(self,node):
        if node.balanceFactor < 0:
            if node.rightChild.balanceFactor > 0:
                self.rotateRight(node.rightChild)
                self.rotateLeft(node)
            else:
                self.rotateLeft(node)

        elif node.balanceFactor > 0:
            if node.leftChild.balanceFactor < 0:
                self.rotateLeft(node.leftChild)
                self.rotateRight(node)
            else:
                self.rotateRight(node)

    def rotateRight(self,rotRoot):
        newRoot = rotRoot.leftChild
        rotRoot.leftChild = newRoot.rightChlid
        if newRoot.RightChild != None:
            newRoot.rightChlid.parent = rotRoot
        newRoot.parent = rotRoot.parent
        if rotRoot.isRoot():    # 跟节点
            self.root = newRoot
        else:
            if rotRoot.isLeftChlid():
                rotRoot.parent.LeftChild = newRoot
            else:
                rotRoot.parent.RightChild = newRoot
        newRoot.rightChlid = rotRoot
        rotRoot.parent = newRoot
        rotRoot.balanceFactor = rotRoot.balanceFactor - 1 + min(newRoot.balanceFactor, 0)
        newRoot.balanceFactor = newRoot.balanceFactor - 1 + max(rotRoot.balanceFactor, 0)

    def rotateLeft(self,rotRoot):
        newRoot = rotRoot.rightChild
        rotRoot.rightChild = newRoot.leftChild
        if newRoot.leftChild != None:
            newRoot.leftChild.parent = rotRoot
        newRoot.parent = rotRoot.parent
        if rotRoot.isRoot():
            self.root = newRoot
        else:   # 搞左右子树
            if rotRoot.isLeftChild():
                rotRoot.parent.leftChild = newRoot
            else:
                rotRoot.parent.rightChild = newRoot
        newRoot.leftChild = rotRoot
        rotRoot.parent = newRoot
        rotRoot.balanceFactor = rotRoot.balanceFactor + 1 - min(newRoot.balanceFactor,0)
        newRoot.balanceFactor = newRoot.balanceFactor + 1 + max(rotRoot.balanceFactor,0)




    # 平衡二叉树删除节点作为练习，可能更加复杂。

mytree = balanceBinarySearchTree()
mytree[3]="red"
mytree[4]="blue"
print(mytree[3])
mytree[6]="yellow"
mytree[2]="at"
mytree[6]="yell"
print(mytree[6])
print(mytree[2])
print(mytree[3])
# mytree.delete(3)
print(mytree[4])
# print(mytree[4].leftChild.payload)    # 这个的测试需要改变返回值为树的节点
